/*************************************************************************
 *
 * This file is part of the SAMRAI distribution.  For full copyright
 * information, see COPYRIGHT and LICENSE.
 *
 * Copyright:     (c) 1997-2020 Lawrence Livermore National Security, LLC
 * Description:   Main program for SAMRAI convection-diffusion ex. problem.
 *
 ************************************************************************/

#include "SAMRAI/SAMRAI_config.h"

#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <fstream>

#ifndef _MSC_VER
#include <unistd.h>
#endif

#include <sys/stat.h>

// Headers for basic SAMRAI objects
#include "SAMRAI/tbox/SAMRAIManager.h"
#include "SAMRAI/tbox/BalancedDepthFirstTree.h"
#include "SAMRAI/mesh/BergerRigoutsos.h"
#include "SAMRAI/pdat/CellData.h"
#include "SAMRAI/tbox/Database.h"
#include "SAMRAI/mesh/TreeLoadBalancer.h"
#include "SAMRAI/tbox/InputDatabase.h"
#include "SAMRAI/tbox/InputManager.h"
#include "MainRestartData.h"
#include "SAMRAI/tbox/SAMRAI_MPI.h"
#include "SAMRAI/hier/Patch.h"
#include "SAMRAI/hier/PatchLevel.h"
#include "SAMRAI/tbox/PIO.h"
#include "SAMRAI/tbox/RestartManager.h"
#include "SAMRAI/tbox/Utilities.h"
#include "SAMRAI/hier/VariableDatabase.h"

// Headers for major algorithm/data structure objects
#include "SAMRAI/geom/CartesianGridGeometry.h"
#include "SAMRAI/geom/CartesianPatchGeometry.h"
#include "SAMRAI/mesh/StandardTagAndInitialize.h"
#include "SAMRAI/algs/MethodOfLinesIntegrator.h"
#include "SAMRAI/hier/PatchHierarchy.h"
#include "SAMRAI/hier/PatchLevel.h"

// Header for application-specific algorithm/data structure object
#include "ConvDiff.h"

// Classes for autotesting.
#if (TESTING == 1)
#include "test/testlib/AutoTester.h"
#endif


#include <vector>
#include <memory>

using namespace SAMRAI;
using namespace algs;

/************************************************************************
 *
 * This is the main program for an AMR application of the method
 * of lines algorithm.  Specifically, it uses a Runge-Kutta method
 * to integrate the convection-diffusion in time with spatial
 * refinement.  The timestep loop in this file controls gridding
 * operations.
 *
 * The gridding objects used in this application are as follows:
 *
 *    hier::PatchHierarchy - A container for the AMR patch hierarchy and
 *       the data on the grid.
 *
 *    geom::CartesianGridGeometry - Defines and maintains the Cartesian
 *       coordinate system on the grid.  The hier::PatchHierarchy
 *       maintains a reference to this object.
 *
 * The Method-of-Lines algorithm uses these two components to integrate
 * data one timestep.  Time-stepping is done in the main routine.
 * This case uses the same timestep across all patches, regardless
 * of the refinement.  See either the Euler or Linear Advection
 * cases for examples of time-refinement cases.
 *
 *    algs::MethodOfLinesIntegrator - advances solution on patches.
 *             Controls integration of quantities over entire patch
 *             hierarchy. Takes the user-defined class that defines
 *             patch operations (ConvDiff in this case) as an argument
 *
 * The user-supplied object defines characteristics and operations to
 * be carried out on each patch.  We pass this object to the integrator
 * which will orchestrate the integration on all patches.
 *
 *    ConvDiff - Defines variables and numerical routines on a single
 *               patch.
 *
 *    mesh::GriddingAlgorithm - Drives the AMR patch hierarchy generation
 *       and regridding procedures.  This object maintains
 *       references to three other algorithmic objects with
 *       which it is configured when they are passed into its
 *       constructor.   They are:
 *
 *       mesh::BergerRigoutsos - Clusters cells tagged for refinement on a
 *          patch level into a collection of logically-rectangular
 *          box domains.
 *
 *       mesh::DistributedLoadBalancer - Processes the boxes generated by the
 *          mesh::BergerRigoutsos algorithm into a configuration from
 *          which patches are contructed.  The algorithm we use in this
 *          class assumes a spatially-uniform workload distribution;
 *          thus, it attempt to produce a collection of boxes
 *          each of which contains the same number of cells.  The
 *          load balancer also assigns patches to processors.
 *
 *       mesh::StandardTagAndInitialize - Couples the gridding algorithm
 *          to the HyperbolicIntegrator. Selects cells for
 *          refinement based on either Gradient detection, Richardson
 *          extrapolation, or pre-defined Refine box region.  The
 *          object maintains a pointer to the algs::MethodOfLinesIntegrator,
 *          which is passed into its constructor, for this purpose.
 *
 ************************************************************************
 */

int main(
   int argc,
   char* argv[])
{
   /*
    * Initialize MPI, SAMRAI, and enable logging.
    */

   tbox::SAMRAI_MPI::init(&argc, &argv);
   tbox::SAMRAIManager::initialize();
   tbox::SAMRAIManager::startup();
   const tbox::SAMRAI_MPI& mpi(tbox::SAMRAI_MPI::getSAMRAIWorld());

   int num_failures = 0;

   /* This extra code block is used to scope some temporaries that are
    * created, it forces the destruction before the manager is
    * shutdown.
    */
   {

      /*
       * Process command line arguments and dump to log file.
       * For non-restarted case, command line is:
       *
       *    executable <input file name>
       *
       * For restarted run, command line is:
       *
       *    executable <input file name> <restart directory> \
       *               <restart number>
       */

      std::string input_filename;
      std::string restart_read_dirname;
      int restore_num = 0;
      bool is_from_restart = false;

      if ((argc != 2) && (argc != 4)) {
         tbox::pout << "USAGE:  " << argv[0] << " <input filename> "
                    << "<restart dir> <restore number> [options]\n"
                    << "  options:\n"
                    << "  none at this time"
                    << std::endl;
         tbox::SAMRAI_MPI::abort();
         return -1;
      } else {
         input_filename = argv[1];
         if (argc == 4) {
            restart_read_dirname = argv[2];
            restore_num = atoi(argv[3]);

            is_from_restart = true;
         }
      }

      tbox::plog << "input_filename = " << input_filename << std::endl;
      tbox::plog << "restart_read_dirname = " << restart_read_dirname << std::endl;
      tbox::plog << "restore_num = " << restore_num << std::endl;

      /*
       * Create input database and parse all data in input file.
       */

      std::shared_ptr<tbox::InputDatabase> input_db(
         new tbox::InputDatabase("input_db"));
      tbox::InputManager::getManager()->parseInputFile(input_filename, input_db);

      /*
       * Retrieve "GlobalInputs" section of the input database and set
       * values accordingly.
       */

      if (input_db->keyExists("GlobalInputs")) {
         std::shared_ptr<tbox::Database> global_db(
            input_db->getDatabase("GlobalInputs"));
         if (global_db->keyExists("call_abort_in_serial_instead_of_exit")) {
            bool flag = global_db->
               getBool("call_abort_in_serial_instead_of_exit");
            tbox::SAMRAI_MPI::setCallAbortInSerialInsteadOfExit(flag);
         }
      }

      /*
       * Retrieve "Main" section of the input database.  First, read dump
       * information, which is used for writing plot files.  Second,
       * if proper restart information was given on command line, and the restart
       * interval is non-zero, create a restart database.
       */

      std::shared_ptr<tbox::Database> main_db(
         input_db->getDatabase("Main"));

      const tbox::Dimension dim(static_cast<unsigned short>(main_db->getInteger("dim")));

      const std::string base_name =
         main_db->getStringWithDefault("base_name", "unnamed");

      const std::string log_filename =
         main_db->getStringWithDefault("log_filename", base_name + ".log");

      bool log_all_nodes = false;
      if (main_db->keyExists("log_all_nodes")) {
         log_all_nodes = main_db->getBool("log_all_nodes");
      }
      if (log_all_nodes) {
         tbox::PIO::logAllNodes(log_filename);
      } else {
         tbox::PIO::logOnlyNodeZero(log_filename);
      }

      int viz_dump_interval = 0;
      if (main_db->keyExists("viz_dump_interval")) {
         viz_dump_interval = main_db->getInteger("viz_dump_interval");
      }

      const std::string visit_dump_dirname =
         main_db->getStringWithDefault("viz_dump_dirname", base_name + ".visit");

      int visit_number_procs_per_file = 1;
      if (viz_dump_interval > 0) {
         if (main_db->keyExists("visit_number_procs_per_file")) {
            visit_number_procs_per_file =
               main_db->getInteger("visit_number_procs_per_file");
         }
      }

      const bool viz_dump_data = (viz_dump_interval > 0);

      int restart_interval = 0;
      if (main_db->keyExists("restart_interval")) {
         restart_interval = main_db->getInteger("restart_interval");
      }

      const std::string restart_write_dirname =
         main_db->getStringWithDefault("restart_write_dirname",
            base_name + ".restart");

#if (TESTING == 1) && !defined(HAVE_HDF5)
      /*
       * If we are autotesting on a system w/o HDF5, the read from
       * restart will result in an error.  We want this to happen
       * for users, so they know there is a problem with the restart,
       * but we don't want it to happen when autotesting.
       */
      is_from_restart = false;
      restart_interval = 0;
#endif

      const bool write_restart = (restart_interval > 0)
         && !(restart_write_dirname.empty());

      /*
       * Get restart manager and root restart database.  If run is from
       * restart, open the restart file.
       */

      tbox::RestartManager* restart_manager = tbox::RestartManager::getManager();

      if (is_from_restart) {
         restart_manager->
         openRestartFile(restart_read_dirname, restore_num,
            mpi.getSize());
      }

      /*
       * Initialize the MainRestartData object which stores the state of the
       * main program for restart.
       */
      MainRestartData* main_restart_data = new MainRestartData(
            "MainRestartData",
            input_db->getDatabase("MainRestartData"));

      /*
       * Create major algorithm and data objects which comprise application.
       * Each object will be initialized either from input data or restart
       * files, or a combination of both.  Refer to each class constructor
       * for details.  For more information on the composition of objects
       * for this application, see comments at top of file.
       */

      std::shared_ptr<geom::CartesianGridGeometry> grid_geometry(
         new geom::CartesianGridGeometry(
            dim,
            "CartesianGeometry",
            input_db->getDatabase("CartesianGeometry")));

      std::shared_ptr<hier::PatchHierarchy> patch_hierarchy(
         new hier::PatchHierarchy(
            "PatchHierarchy",
            grid_geometry,
            input_db->getDatabase("PatchHierarchy")));

      ConvDiff* convdiff_model = new ConvDiff("ConvDiff",
            dim,
            input_db->getDatabase("ConvDiff"),
            grid_geometry);

      std::shared_ptr<algs::MethodOfLinesIntegrator> mol_integrator(
         new algs::MethodOfLinesIntegrator(
            "MethodOfLinesIntegrator",
            input_db->getDatabase("MethodOfLinesIntegrator"),
            convdiff_model));

      std::shared_ptr<mesh::StandardTagAndInitialize> error_detector(
         new mesh::StandardTagAndInitialize(
            "StandardTagAndInitialize",
            mol_integrator.get(),
            input_db->getDatabase("StandardTagAndInitialize")));

      std::shared_ptr<mesh::BergerRigoutsos> box_generator(
         new mesh::BergerRigoutsos(
            dim,
            input_db->getDatabaseWithDefault(
               "BergerRigoutsos",
               std::shared_ptr<tbox::Database>())));

      std::shared_ptr<mesh::TreeLoadBalancer> load_balancer(
         new mesh::TreeLoadBalancer(
            dim,
            "LoadBalancer",
            input_db->getDatabase("LoadBalancer"),
            std::shared_ptr<tbox::RankTreeStrategy>(new tbox::BalancedDepthFirstTree)));
      load_balancer->setSAMRAI_MPI(tbox::SAMRAI_MPI::getSAMRAIWorld());

      std::shared_ptr<mesh::GriddingAlgorithm> gridding_algorithm(
         new mesh::GriddingAlgorithm(
            patch_hierarchy,
            "GriddingAlgorithm",
            input_db->getDatabase("GriddingAlgorithm"),
            error_detector,
            box_generator,
            load_balancer));

      /*
       * Set up Visualization plot file writer(s).
       */
#ifdef HAVE_HDF5
      std::shared_ptr<appu::VisItDataWriter> visit_data_writer(
         new appu::VisItDataWriter(
            dim,
            "ConvDiff VisIt Writer",
            visit_dump_dirname,
            visit_number_procs_per_file));
      convdiff_model->registerVisItDataWriter(visit_data_writer);
#endif

      /*
       * After creating all objects and initializing their state, we
       * print the input database and variable database contents to
       * the log file.
       */

      tbox::plog << "\nCheck input data and variables before simulation:"
                 << std::endl;
      tbox::plog << "Input database..." << std::endl;
      input_db->printClassData(tbox::plog);
      tbox::plog << "\nVariable database..." << std::endl;
      hier::VariableDatabase::getDatabase()->printClassData(tbox::plog);
      mol_integrator->initializeIntegrator(gridding_algorithm);

      /****************************************************************
      *
      *  INITIALIZE DATA ON PATCHES
      *
      ****************************************************************
      *
      *  Build patch hierarchy and initialize the data on the patches
      *  in the hierarchy. Note: this step is performed by the
      *  algs::TimeRefinementIntegrator in Euler/LinAdv example cases.
      *  1) Create a "tag_buffer" for each level in the Hierarchy.
      *  2) Create the coarse (i.e. level 0) grid.
      *  3) Cycle through levels 1-max_levels, initializing data
      *     on each.  The makeFinerLevel method calls the error
      *     estimator (remember, it was registered with the
      *     gridding algorithm object) and tags cells for refinement
      *     as it generates patches on the finer levels.
      *  4) Dump this initial data to viz file.
      *
      ****************************************************************/

      std::vector<int>
      tag_buffer_array(patch_hierarchy->getMaxNumberOfLevels());
      for (int il = 0; il < patch_hierarchy->getMaxNumberOfLevels(); ++il) {
         tag_buffer_array[il] = main_restart_data->getTagBuffer();
         tbox::pout << "il = " << il << " tag_buffer = "
                    << tag_buffer_array[il]
                    << std::endl;
      }

      std::vector<double> regrid_start_time(
         patch_hierarchy->getMaxNumberOfLevels());

      double loop_time = main_restart_data->getLoopTime();
      int loop_cycle = main_restart_data->getIterationNumber();

      if (tbox::RestartManager::getManager()->isFromRestart()) {

         patch_hierarchy->initializeHierarchy();

         gridding_algorithm->getTagAndInitializeStrategy()->
         resetHierarchyConfiguration(patch_hierarchy,
            0,
            patch_hierarchy->getFinestLevelNumber());

      } else {

         gridding_algorithm->makeCoarsestLevel(loop_time);

         bool done = false;
         bool initial_cycle = true;
         for (int ln = 0;
              patch_hierarchy->levelCanBeRefined(ln) && !done;
              ++ln) {
            gridding_algorithm->makeFinerLevel(
               tag_buffer_array[ln],
               initial_cycle,
               loop_cycle,
               loop_time);
            done = !(patch_hierarchy->finerLevelExists(ln));
         }
      }

      tbox::RestartManager::getManager()->closeRestartFile();

#if (TESTING == 1)
      /*
       * Create the autotesting component which will verify correctness
       * of the problem. If no automated testing is done, the object does
       * not get used.
       */
      AutoTester* autotester = new AutoTester("AutoTester", dim, input_db);
#endif

      /*******************************************************************
       *
       *  MAIN TIME ADVANCE LOOP
       *
       *******************************************************************
       *
       *  1) Set start and end time.
       *  2) Start integration timesteps.
       *     While (loop_time < end_time) {
       *        2a) Write restart data.
       *        2b) Advance all levels in the hierarchy by time
       *            dt by calling algs::MethodOfLinesIntegrator's
       *            advanceHierarchy method.
       *        2c) Check if it is time to do a regrid step.  If so,
       *            have the mesh::GriddingAlgorithm
       *            call its regridAllFiner
       *            Levels method, passing in the coarsest (0) level.
       *            This method will invoke the Gradient detector,
       *            Berger Rigoutsos algorithm, and load balancer
       *            while generating finer grid levels.
       *     }
       *
       ******************************************************************/
      int iteration_num = main_restart_data->getIterationNumber();

#if (TESTING == 1)
      /*
       * If we are doing autotests, check result...
       */
      num_failures += autotester->evalTestData(iteration_num,
            patch_hierarchy,
            loop_time,
            mol_integrator,
            gridding_algorithm);
#endif

      if (viz_dump_data) {

#ifdef HAVE_HDF5
         visit_data_writer->writePlotData(
            patch_hierarchy,
            iteration_num,
            loop_time);

#endif
      }

      while ((loop_time < main_restart_data->getEndTime()) &&
             (iteration_num < main_restart_data->getMaxTimesteps())) {

         iteration_num = main_restart_data->getIterationNumber();
         ++iteration_num;

         tbox::pout << "++++++++++++++++++++++++++++++++++++++++++++" << std::endl;
         tbox::pout << "At begining of timestep # " << iteration_num - 1
                    << std::endl;
         tbox::pout << "Simulation time is " << loop_time << std::endl;

         double dt = mol_integrator->getTimestep(patch_hierarchy, loop_time);

         mol_integrator->advanceHierarchy(patch_hierarchy, loop_time, dt);

         loop_time += dt;

         tbox::pout << "At end of timestep # " << iteration_num - 1 << std::endl;
         tbox::pout << "Simulation time is " << loop_time << std::endl;
         tbox::pout << "++++++++++++++++++++++++++++++++++++++++++++" << std::endl;

         /*
          * Write restart file at specified intervals.  Set current state
          * of "main" in the main_restart_data object, before writing restart.
          */
         main_restart_data->setLoopTime(loop_time);
         main_restart_data->setIterationNumber(iteration_num);
         if (write_restart) {
            if ((iteration_num % restart_interval) == 0) {
               tbox::RestartManager::getManager()->
               writeRestartFile(restart_write_dirname,
                  iteration_num);
            }
         }

         /*
          * At specified intervals, write out data files for plotting.
          */
#ifdef HAVE_HDF5
         if (viz_dump_data) {
            if ((iteration_num % viz_dump_interval) == 0) {

               visit_data_writer->writePlotData(patch_hierarchy,
                  iteration_num,
                  loop_time);
            }
         }
#endif

         /*
          *  At specified intervals, regrid.
          */
         if ((iteration_num % main_restart_data->getRegridStep()) == 0 &&
             patch_hierarchy->getMaxNumberOfLevels() > 1) {
            tbox::plog << "\n\n############################################"
                       << std::endl;
            tbox::plog << "                 REGRIDDING" << std::endl;
            tbox::plog << "Finest level before regrid: "
                       << patch_hierarchy->getFinestLevelNumber() << std::endl;

            gridding_algorithm->regridAllFinerLevels(0,
               tag_buffer_array,
               iteration_num,
               loop_time,
               regrid_start_time);

            tbox::plog << "Finest level after regrid: "
                       << patch_hierarchy->getFinestLevelNumber() << std::endl;
            tbox::plog << "############################################\n\n"
                       << std::endl;
         }

#if (TESTING == 1)
         /*
          * If we are doing autotests, check result...
          */
         num_failures += autotester->evalTestData(iteration_num,
               patch_hierarchy,
               loop_time,
               mol_integrator,
               gridding_algorithm);
#endif

      }

      /*
       * At conclusion of simulation, deallocate objects.
       */

      box_generator.reset();

      load_balancer.reset();
      gridding_algorithm.reset();

#ifdef HAVE_HDF5
      visit_data_writer.reset();
#endif

      //delete tag_buffer_array;

      error_detector.reset();
      mol_integrator.reset();

      if (convdiff_model) delete convdiff_model;

      patch_hierarchy.reset();

      grid_geometry.reset();

      if (main_restart_data) delete main_restart_data;

      main_db.reset();
      input_db.reset();

#if (TESTING == 1)
      delete autotester;
#endif

   }

   if (num_failures == 0) {
      tbox::pout << "\nPASSED:  ConvDiff" << std::endl;
   }

   tbox::SAMRAIManager::shutdown();
   tbox::SAMRAIManager::finalize();
   tbox::SAMRAI_MPI::finalize();

   return num_failures;
}
